﻿module EmitOptimizer

open Ast
open EmitType
open StateMonad
open SymbolTableTypes
open System.Collections.Generic

type private Environment = 
  { 
    Labels : List<List<EmitLabel>>    
  }


type private State<'a> = Environment -> ('a * Environment)

let rec private optimize_emit_code (code : seq<Emit>) : State<Unit> =
  state{
    if code |> Seq.length >= 2 then
      let! state = getState
      let is_label = (fun line -> match line with |Emit.Label _ -> true | _ -> false)
      let labels = code |> Seq.takeWhile is_label
      if labels |> Seq.length > 1 then
        let labels = labels |> Seq.map (fun (Emit.Label name) -> name)
        do state.Labels.Add(List(labels))

      return! optimize_emit_code (code |> Seq.skipWhile is_label |> Seq.skip 1)
    else
      return ()
  }

and private optimize_code code =
  state{
    let! state = getState
    let to_remove = List<Emit>()
    do for l in code do
        match l with
        | Emit.Jmp (Fjmp (_, label))
        | Emit.Jmp (Tjmp (_, label))
        | Emit.Jmp (EmitJmp.Jmp label) ->
          let names = state.Labels |> Seq.tryFind (fun s -> s.Contains label)
          match names with
          | None -> ()
          | Some names -> 
            do label := !names.[0]
        | Emit.Label label as lab ->
          let names = state.Labels |> Seq.tryFind (fun s -> s.Contains label)
          if names.IsSome && names.Value.[0] <> label then
            do to_remove.Add(lab)
        | _ -> ()
    let opt = code |> Seq.filter (fun l -> l |> to_remove.Contains |> not) |> Seq.toList

    return opt
  }

let optimize_emit code = 
  let state = { Labels = List() }
  let is_not_comment = (fun line -> match line with | Emit.Comment _ -> false | _ -> true)
  let env = execute (optimize_emit_code (code |> Seq.filter is_not_comment)) state |> snd
  let opt = execute (optimize_code code) env |> fst
  ()